import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import LSTM, Dense, Dropout
import matplotlib.pyplot as plt
import os
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error
import math

maotai = pd.read_csv('SH601229.csv')

# 前2126天开盘价作为训练集，表格从0开始计数，2:3是提取2-3列，前闭后开，所以是提取c列开盘价
training_set = maotai.iloc[0:1184 - 100, 2:3].values
test_set = maotai.iloc[1184 - 100:, 2:3].values

# 归一化
sc = MinMaxScaler(feature_range=(0, 1))
training_set_scaled = sc.fit_transform(training_set)
test_set = sc.fit_transform(test_set)

x_train, y_train = [], []
x_test, y_test = [], []

# 前60天作为输入数据，第60天作为标签
for i in range(50, len(training_set_scaled)):
    x_train.append(training_set_scaled[i - 50:i, 0])
    y_train.append(training_set_scaled[i, 0])

np.random.seed(8)
np.random.shuffle(x_train)
np.random.seed(8)
np.random.shuffle(y_train)
tf.random.set_seed(8)

# 变换数据格式
x_train, y_train = np.array(x_train), np.array(y_train)

# 使得输入数据符合RNN要求
x_train = np.reshape(x_train, (x_train.shape[0], 50, 1))

for i in range(50,len(test_set)):
    x_test.append(test_set[i-50:i,0])
    y_test.append(test_set[i,0])

x_test, y_test = np.array(x_test), np.array(y_test)
x_test = np.reshape(x_test, (x_test.shape[0], 50,1))

# 定义网络
model = tf.keras.Sequential([
    LSTM(80,return_sequences=True),
    Dropout(0.2),
    LSTM(100),
    Dropout(0.2),
    Dense(1)
])

# 只观测loss的数值，不观测准确率，所以删去metrics选项，一会在每个epoch迭代显示时只显示loss值
model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
              loss='mean_squared_error') # 损失函数用均方误差
checkpoint_save_path = "./28 checkpoint/stock_predict.ckpt"

if os.path.exists(checkpoint_save_path + '.index'):
    print('-------------load the model-----------------')
    model.load_weights(checkpoint_save_path)

else:
    cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
                                                     save_weights_only=True,
                                                     save_best_only=True,
                                                     monitor='loss')  # 由于fit没有给出测试集，不计算测试集准确率，根据loss，保存最优模型

    model.fit(x_train,y_train,batch_size=32,epochs=100,validation_data=(x_test,y_test),validation_freq=1,
              callbacks=[cp_callback])

    model.summary()


# 预测
predict = model.predict(x_test)
predict = sc.inverse_transform(predict)
real = sc.inverse_transform(test_set[50:])
# 画出真实数据和预测数据的对比曲线
plt.plot(real,color='red',label='stock price')
plt.plot(predict,color='green',label = 'predict price')
plt.title('stock prediction')
plt.xlabel('time')
plt.ylabel('price')
plt.legend()
plt.show()